#!/usr/bin/perl

# mkv4ps3
#
# Copyright Josh Carroll <josh.carroll@gmail.com>
#
# Released under the GPL v2
#
# This script takes an MKV file as input and produces either
#	an AVI (xvid/ac3) or MP4 (h264/aac) which will play back
#	on the Sony PlayStation 3 (tm).
#
# By default, if the input file is 720p (or less), the script
#	will retain the h264 elemental stream and use it unmodified,
#	while transcoding the ac3 to 5.1 aac with proper channel
#	mapping.

# If the -target option is used to specify avi, then
#	the script uses ffmpeg to transcode to a high bitrate XVID
#	avi file, leaving the AC3 audio intact.
#
# If the input file is 1080p, then the script prompts the user
#	whether they would like to transcode the h264 elemental stream
#	to a high quality h264 stream that's Level High@4.1 compliant and
#	mux that with the transcoded aac audio, or whether they want to
#	transcode to avi (xvid/ac3). If the user has explicitly specified
#	the target, they are not prompted.

use warnings;
use strict;
use Getopt::Long;
use POSIX qw(floor);
use Cwd qw(realpath);

my $prog = $0;
$prog =~ s/.*\///img;

my ($opt_in, $opt_out, $opt_tmp, $opt_nosplit, $opt_thr, $opt_quiet, $opt_reenc,
	$opt_help, $opt_target, $opt_pad, $opt_2pass, $opt_verbose, $opt_force);

GetOptions( 'in=s' => \$opt_in,
			'out=s' => \$opt_out,
			'tmp=s' => \$opt_tmp,
			'nosplit' => \$opt_nosplit,
			'thr=s' => \$opt_thr,
			'target=s' => \$opt_target,
			'pad=s' => \$opt_pad,
			'2pass' => \$opt_2pass,
			'force' => \$opt_force,
			'h' => \$opt_help,
			'v' => \$opt_verbose);

if(! $opt_in || ! $opt_out || $opt_help) {
	print "Usage: $prog -in <in.mkv> -out <out dir>\n";
	print "		[ -tmp <dir> | -nosplit | -thr N | \n";
	print "			-target avi|mp4 | -pad <kbps>]\n\n";
	print "The -target option is generally only needed for 1080p input\n";
	print "	unless you want an avi (xvid/ac3), since the default is h264/aac\n\n";
	print "If target mode is avi and you want to increase the video bitrate,\n";
	print "	add -pad <kbps> to add additional bitrate to the xvid video.\n\n";
	exit 0;
}

# setup the tmp dir if necessary
setup_working_dirs();

# assume we need roughly 2x the input file size for storage (worse case scenario)
check_available_diskspace($opt_in);

check_required_tools();
tprint("Checking input file sanity...\n");
check_resolution($opt_in);

# output file(s) are based on input file name
if(! $opt_target) {
	if($opt_out !~ /\.mp4/i) {
		print "Warning: default target is avi but output file has a different extension.\n";
	}
} elsif( ($opt_target eq "mp4" && $opt_out !~ /\.mp4/i ) ||
	($opt_target eq "avi" && $opt_out !~ /\.avi/i) ) {
	print "Warning: target specified is $opt_target but output file has a different extension.\n";
}

if(defined $opt_target && $opt_target ne "mp4" && $opt_target ne "avi") {
	print "Invalid target: $opt_target\n";
	exit 1;
}

# setup tmp files
my $tmp_ac3 = "$opt_tmp/tmp.ac3";
my $tmp_dts = "$opt_tmp/tmp.dts";
my $tmp_h264 = "$opt_tmp/tmp.h264";
my $tmp_pre_reenc_h264 = "$opt_tmp/tmp_reenc.h264";
my $tmp_wav = "$opt_tmp/tmp.wav";
my $tmp_wav_fifo = "$opt_tmp/tmp.fifo.wav";
my $tmp_aac = "$opt_tmp/tmp.aac";
my $tmp_mp4 = "$opt_tmp/tmp.mp4";
my $tmp_2pass_log = "$opt_tmp/ffmpeg2pass";
my $tmp_2pass_real_path = "$tmp_2pass_log-0.log";

# if the user wants to transcode directly to avi, we skip all of the many
#	intermediate steps for mp4 creation (or re-encoding)
my $dont_extract = 0;
if(defined $opt_target && $opt_target eq "avi") {
	$dont_extract = 1;
	tprint("Reading mkv info...\n");
} else {
	tprint("Extracting mkv file...\n");
}

# get the fps of the video and the duration (in seconds) of the MKV, and
#	extract the mkv (if only re-muxing to mp4 with aac)
my ($vid_fps,$duration,$audio_file) = extract_mkv($opt_in, $tmp_h264, $dont_extract);

if(defined $opt_target && $opt_target eq "avi") {
	tprint("Transcoding mkv to avi (this may take a while)...\n");
	mkv2avi($opt_in, $opt_out, $vid_fps, $duration, $tmp_2pass_log);
	tprint("Done!\n");
	unlink($tmp_2pass_real_path);
	rmdir($opt_tmp);
	exit(0);
} elsif(defined $opt_target && $opt_target eq "mp4") {
	tprint("Re-encoding h.264 elementa stream (this may take a while)...\n");
	reencode_h264($tmp_h264, $tmp_pre_reenc_h264, $duration, $tmp_2pass_log);
} else {
	tprint("Changing h.264 stream's level_idc header...\n");
	change_h264_level($tmp_h264);
}

tprint("Converting audio to aac...\n");
conv_to_aac($audio_file, $tmp_aac);

tprint("Creating mp4 file...\n");
create_mp4($tmp_h264, $tmp_aac, $vid_fps, $tmp_mp4, $opt_out);

# cleanup
tprint("Cleaning up tmp directory...\n");
unlink($tmp_ac3,$tmp_h264,$tmp_wav,$tmp_aac,$tmp_mp4,$tmp_pre_reenc_h264,$tmp_2pass_real_path);
rmdir($opt_tmp);

tprint("Done!\n");
# done


sub check_resolution {
	my $mkv = shift;

	if($opt_verbose) {
		print "Running: mkvinfo \"$mkv\" 2>&1\n";
	}
	my $output = `mkvinfo "$mkv" 2>&1`;
	chomp($output);

	my $width = 0;
	if($output =~ /Pixel width:\s+(\d+)/i) {
		$width = $1;
	}
	if($width == 0) {
		print "Could not determine video width. mkvinfo parsing failed. mkvinfo output:\n\n$output\n";
		exit 1;
	}

	my $height = 0;
	if($output =~ /Pixel height:\s+(\d+)/i) {
		$height = $1;
	}
	if($height == 0) {
		print "Could not determine video width. mkvinfo parsing failed. mkvinfo output:\n\n$output\n";
		exit 1;
	}

	if($width > 1280) {
		if(! $opt_target) {
			print "The input file appears to need transcoding.\n";
			print "Would you like to transcode to avi (xvid/ac3: faster,\n";
			print "but slightly less video quality), or mp4 (h264/aac: slower,\n";
			print "and better video quality, but potential some loss in\n";
			print "audio quality).\n\n";
			
			while(! $opt_target) {
				print "Please enter your choice (avi,mp4): ";
				my $answer = <STDIN>;
				chomp($answer);
				if($answer =~ /avi/i) {
					$opt_target = "avi";
				} elsif($answer =~ /mp4/i) {
					$opt_target = "mp4";
				} else {
					print "Invalid response! Please try again\n";
				}
			}
		}
	}
	
	if($height % 16) {
		print "Error, height must be a multiple of 16 for the PS3 to recognize the file.\n";
		print "Consider using ffmpeg/mencoder to crop this video first.\n";
		exit 1;
	}
}


sub extract_mkv {
	my ($in, $h264, $dont_extract) = @_;

	my $fps = -1;
	my $duration = -1;
	my $type = undef;
	my $audio_type = undef;

	my ($audio_track_num, $h264_track_num) = (-1, -1);
	my ($found_video, $found_audio) = (0,0);

	# find out which track is which
	my $curr_track_num = -1;
	if($opt_verbose) {
		print "Running: mkvinfo \"$in\" 2>&1\n";
	}
	my @output = `mkvinfo "$in" 2>&1`;
	chomp(@output);
	foreach my $line (@output) {
		if($found_video && $found_audio && $fps != -1 && defined $audio_type) {
			last;
		}
		if($line =~ /Default duration.*\((\d+\.\d+)\s*fps/i) {
			my $tmp = $1;
			if($type =~ /video/i) {
				$fps = $tmp;
			}
		} elsif($line =~ /^\s*\|\s*\+\s*Duration:\s*(\d+\.\d+)\s*s\s+/i) {
			$duration = $1;
		} elsif($line =~ /Track number:\s+(\d+)/i) {
			$curr_track_num = $1;
			next;
		} elsif($line =~ /Track type:\s+(\S+)/i) {
			if($curr_track_num == -1) {
				# we haven't found the track # yet, but we found the type. this shouldn't happen
				print "Error, mkvinfo failed.\n";
				exit 1;
			}
			$type = $1;
			if($type =~ /audio/i) {
				if(! $found_audio) {
					$audio_track_num = $curr_track_num;
					$found_audio = 1;
				}
			} else {
				if(! $found_video) {
					$h264_track_num = $curr_track_num;
					$found_video = 1;
				}
			}
		} elsif($found_audio && $line =~ /Codec ID:\s+(\S+)/i) {
			if(! $audio_type) {
				$audio_type = $1;
			}
		}	
	}

	if($fps == -1 || $duration == -1 || $audio_track_num == -1 || $h264_track_num == -1) {
		print "mkvinfo parsing failed. mkvinfo output:\n\n" . join("\n", @output), "\n";
		exit 1;
	}

	my $audio_file;
	if($audio_type eq "A_DTS") {
		$audio_file = $tmp_dts;
	} elsif($audio_type eq "A_AC3") {
		$audio_file = $tmp_ac3;
	} else {
		print "Error: invalid audio type in mkv: $audio_type\n";
		exit 1;
	}

	if(! $dont_extract) {
		if($opt_verbose) {
			print "Running: mkvextract tracks \"$in\" $audio_track_num:$audio_file $h264_track_num:$tmp_h264 2>&1\n";
		}
		my $output = `mkvextract tracks "$in" $audio_track_num:$audio_file $h264_track_num:$tmp_h264 2>&1`;
		if($?) {
			print "Error, mkvextract failed! output:\n\n$output\n";
			exit 1;
		}
	}

	if($fps == -1) {
		print "Error, could not determine input MKV fps\n";
		exit 1;
	}

	return ($fps, $duration, $audio_file);
}


sub is_dir_empty {
	my $dir = shift;

	my $count = 0;
	opendir my $fh, "$dir" or die "Cannot open directory: $dir\n";
	while(my $ent = readdir($fh)) {
		if($ent eq "." || $ent eq "..") {
			next;
		}

		$count++;
	}

	if($count > 0) {
		return 0;
	} else {
		return 1;
	}
}


sub change_h264_level {
	my $file = shift;

	open my $fh, "+< $file" or die "Cannot open temporary h264 file\n";
	sysseek($fh, 7, 0);
	syswrite($fh,"\x29",1);
	close($fh);

}

sub conv_to_aac {
	my ($in, $aac) = @_;

	my $fifo = $tmp_wav_fifo;

	if($opt_verbose) {
		print "Running mkfifo $fifo\n";
	}
	`mkfifo $fifo`;
	if($?) {
		print "mkfifo failed -- $!\n";
		exit 1;
	}

	if($opt_verbose) {
		print "Running neroAacEnc -ignorelength -q 1.0 -if $fifo -of $aac > /dev/null 2>&1 &\n";
	}
	system("neroAacEnc -ignorelength -q 1.0 -if $fifo -of $aac > /dev/null 2>&1 &");

	if($opt_verbose) {
		print "Running mplayer $in -af channels=6:6:0:0:1:1:2:4:3:5:4:2:5:3 -ao pcm:fast:waveheader:file=$fifo -channels 6 -novideo 2>&1\n";
	}
	my $output = `mplayer $in -af channels=6:6:0:0:1:1:2:4:3:5:4:2:5:3 -ao pcm:fast:waveheader:file=$fifo -channels 6 -novideo 2>&1`;
	if($?) {
		tprint("mplayer/nero conversion to aac failed: $output\n");
		exit 1;
	}

	# clean up the fifo and the ac3/dts file to save on disk space while we're muxing
	unlink($in);
	unlink($fifo);
}

sub create_mp4 {
	my ($h264, $aac, $fps, $tmp_mp4, $outf) = @_;

	# size of h264 + aac
	my $total_size = 0;

	# split the file into 3.5 G chunks to avoid problems with PS3 playback of 4G files
	# this is the size in KB
	my $split_size = 4000000;

	my $split = 0;

	if(! $opt_nosplit) {
		# find out how big the existing file is. If it's less than our split size, then
		#	don't bother splitting it
		my ($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
			$atime,$mtime,$ctime,$blksize,$blocks) = stat($h264);

		$total_size += $size;

		($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
			$atime,$mtime,$ctime,$blksize,$blocks) = stat($aac);

		$total_size += $size;

		# normalize to KB for comparison
		$total_size /= 1024;
		if($total_size > $split_size) {
			$split = 1;
		}
	}

	my $split_str = "";
	if($split) {
		$split_str = "-splits $split_size";
	} else {
		$split_str = "";
	}

	my $cmd = "mp4box $split_str -tmp $opt_tmp -add $h264 -fps $fps -add $aac -new \"$outf\" 2>&1";
	if($opt_verbose) {
		print "Running $cmd\n";
	}
	my $output = `$cmd`;
	if($?) {
		print "mp4box creation of mp4 failed! output:\n\n$output\n";
		exit 1;
	}
}

sub check_required_tools {

	my $errors = 0;

	my @required = qw(mkvinfo mkvextract ffmpeg mplayer mp4box neroAacEnc);
	foreach my $req_prog (@required) {
		my $found = 0;
		foreach my $path_ent (split(/:/, $ENV{PATH})) {
			if(-r "$path_ent/$req_prog") {
				$found = 1;
				last;
			}
		}
		if(! $found) {
			$errors++;
			print "Error, required program: $req_prog not found. Please install it.\n";
		}
	}

	if($errors > 0) {
		exit 1;
	}
}

sub tprint {
	my $msg = shift;

	if($opt_quiet) {
		return;
	}

	print $msg;
}

sub reencode_h264 {
	my ($in, $out, $duration, $tmp_2pass_log) = @_;

	# we want the file for the next step to be in the path pointed
	#	to be $in, so move it
	`mv -f $in $out`;

	# determine the input bitrate using size of file and duration
	my ($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
		$atime,$mtime,$ctime,$blksize,$blocks) = stat($out);

	# size is in bytes, we want bits
	$size *= 8;

	# figure out bitrate in kbps
	my $vbitrate = POSIX::floor($size / $duration / 1024);

	# in/out are "swapped" here, since we want the resulting file to be the
	#	$tmp_h264 path
	my $thr_str = "";
	if($opt_thr) {
		$thr_str = "-threads $opt_thr";
	}

	my $cmd = "ffmpeg -y $thr_str %PASS% -f h264 -i $out -vcodec libx264 -level 41 -coder 1 -flags +loop -flags2 +dct8x8 -cmp +chroma -partitions +parti4x4+partp8x8+partb8x8 -me umh -subq 5 -me_range 16 -g 250 -keyint_min 25 -sc_threshold 40 -i_qfactor 0.71 -rc_eq 'blurCplx^(1-qComp)' -qcomp 0.6 -qmin 10 -qmax 51 -qdiff 4 -refs 3 -bf 3 -trellis 1 -b ${vbitrate}k -maxrate ${vbitrate}k -f h264 %OUT%";

	if($opt_2pass) {
		print "Using 2-pass encode...\n";
		my $cmd1 = $cmd;
		$cmd1 =~ s/%PASS%/-pass 1 -passlogfile $tmp_2pass_log/ig;
		$cmd1 =~ s/%OUT%/\/dev\/null/ig;

		if($opt_verbose) {
			print "Running (1st pass) $cmd1 2>&1\n";
		}
		my $pass1_out = `$cmd1 2>&1`;
		if($?) {
			print "First pass failed for x264 re-encode! output:\n\n$pass1_out\n";
			exit 1;
		}

		my $cmd2 = $cmd;
		$cmd1 =~ s/%PASS%/-pass 2 -passlogfile $tmp_2pass_log/ig;
		$cmd2 =~ s/%OUT%/$out/ig;

		if($opt_verbose) {
			print "Running (2nd pass) $cmd2 2>&1\n";
		}
		my $pass2_out = `$cmd2 2>&1`;
		if($?) {
			print "Second pass failed for x264 re-encode! output:\n\n$pass2_out\n";
			exit 1;
		}
	} else {
		my $cmd1 = $cmd;
		$cmd1 =~ s/%PASS%//ig;
		$cmd1 =~ s/%OUT%/$out/ig;

		if($opt_verbose) {
			print "Running $cmd1 2>&1\n";
		}
		my $output = `$cmd1 2>&1`;
		if($?) {
			print "Re-encode of h264 stream failed. output:\n\n$output\n";
			exit 1;
		}
	}
}

sub mkv2avi {
	my ($in, $outf, $fps, $duration, $tmp_2pass_log) = @_;

	my $vbitrate = get_video_bitrate($in);

	if($opt_pad) {
		$vbitrate += $opt_pad;
	}

	my $thr_str = "";
	if($opt_thr) {
		$thr_str = "-threads $opt_thr";
	}
	
	my $cmd = "ffmpeg -y $thr_str %PASS% -i $in -vcodec mpeg4 -vtag XVID -b ${vbitrate}k -cmp 2 -subcmp 2 -trell 1 -mbd 2 -acodec copy -r $fps -f avi %OUT%";

	if($opt_2pass) {
		print "Using 2-pass encode...\n";
		my $cmd1 = $cmd;
		$cmd1 =~ s/%PASS%/-pass 1 -passlogfile $tmp_2pass_log/ig;
		$cmd1 =~ s/%OUT%/\/dev\/null/ig;

		if($opt_verbose) {
			print "Running (1st pass) $cmd1 2>&1\n";
		}
		my $pass1_out = `$cmd1 2>&1`;
		if($?) {
			print "First pass failed for mkv -> avi conversion! output:\n\n$pass1_out\n";
			exit 1;
		}

		my $cmd2 = $cmd;
		$cmd2 =~ s/%PASS%/-pass 2 -passlogfile $tmp_2pass_log/ig;
		$cmd2 =~ s/%OUT%/$outf/ig;

		if($opt_verbose) {
			print "Running (2nd pass) $cmd2 2>&1\n";
		}
		my $pass2_out = `$cmd2 2>&1`;
		if($?) {
			print "Second pass failed for mkv -> avi conversion! output:\n\n$pass2_out\n";
			exit 1;
		}
	} else {
		my $cmd1 = $cmd;
		$cmd1 =~ s/%PASS%//ig;
		$cmd1 =~ s/%OUT%/$outf/ig;

		if($opt_verbose) {
			print "Running $cmd1 2>&1\n";
		}
		my $output = `$cmd1 2>&1`;
		if($?) {
			print "mkv to avi conversion failed. output:\n\n$output\n";
			exit 1;
		}
	}
}


sub get_video_bitrate {
	my $input = shift;

	my $bitrate = 0;

	if($opt_verbose) {
		print "Running ffmpeg -i \"$input\" 2>&1\n";
	}
	my $info = `ffmpeg -i "$input" 2>&1`;

	# determine the video bitrate, either from ffmpeg -i output directly
	#   or by assuming a constant average bitrate over the length of the video
	if($info =~ /Duration.*bitrate:\s+(\d+)\s+kb\/s/i) {
		$bitrate = $1;
	} else {
		# need to find the file size and video length to guesstimate
		my ($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
			$atime,$mtime,$ctime,$blksize,$blocks)
			= stat($input);
		$info =~ /Duration:\s+(\d+):(\d+):(\d+\.\d+),/img;
		if(! $1) {
			return 0;	
		}
		my $seconds = $1*60*60 + $2*60 + $3;
		$bitrate = floor($size*8 / $seconds / 1024);
	}

	if($bitrate == 0) {
		print "Could not determine bitrate! Corrupt file?\n";
		exit 1;
	}

	return $bitrate;
}

sub setup_working_dirs {

	my $default_tmp_dir = "/tmp/$prog";
	if(! $opt_tmp) {
		print "Warning, using default tmp directory: $default_tmp_dir\n";
		$opt_tmp = $default_tmp_dir;
	} else {
		$opt_tmp = realpath($opt_tmp);
	}

	if(-f $opt_tmp) {
		print "Error, tmp directory is already a file.\n";
		exit 1;
	} elsif(-d $opt_tmp) {
		if(! is_dir_empty($opt_tmp)) {
			print "Error, tmp directory is not empty.\n";
			exit 1;
		}
	} else {
		mkdir($opt_tmp, 0755);
		if(! -d $opt_tmp) {
			print "Could not create tmp directory!\n";
			exit 1;
		}
	}

	if(-f $opt_out) {
		print "Error, requested output file already exists!\n";
		exit 1;
	} elsif(-d $opt_out) {
		print "Error, requested output file is a directory.\n";
		exit 1;
	}
}

sub check_available_diskspace {
	my $in = shift;

	my $needed_size = 0;
	my $available_size = 0;
	my ($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
		$atime,$mtime,$ctime,$blksize,$blocks) = stat($in);

	# rough size needed (in bytes)
	$needed_size = 2 * $size;

	# check what's available in the specified tmp dir
	my $df = `BLOCKSIZE=1024 df $opt_tmp`;
	chomp($df);

	$df =~ /^\/\S+\s+(\d+)/imgs;
	$available_size = $1 * 1024;

	if($needed_size > $available_size) {
		if($opt_force) {
			print "Warning: temp directory may not have enough free space. But -force specified...\n";
		} else {
			my $human_size = $needed_size / 1024 / 1024;
			print "Error: not enough disk space to continue.\n";
			printf("Please use -tmp to specify a temp directory with at least %.2f MB free\n", $human_size);
			exit 1;
		}
	}
}
